import torch
import torch.nn as nn
import torch.nn.functional as F


class Attention(nn.Module):
    def __init__(self, query_size, key_size, value_size1, value_size2, output_size):
        """
        :param query_size: Q的最后一个维度
        :param key_size: K的最后一个维度
        :param value_size1: V的尺寸为(1, value_size1, value_size2)
        :param value_size2:
        :param output_size:
        """
        super(Attention, self).__init__()
        self.query_size = query_size
        self.key_size = key_size
        self.value_size1 = value_size1
        self.value_size2 = value_size2
        self.output_size = output_size

        # 初始化注意力机制实现中第一步的线性层
        self.attn = nn.Linear(self.query_size + self.key_size, value_size1)

        # 初始化注意力机制实现中第三步的线性层
        self.attn_combine = nn.Linear(self.query_size + self.value_size2, self.output_size)

    def forward(self, Q, K, V):
        # 假定QKV都是三维张量
        # 第1步，将QK进行纵轴的拼接，然后做一次线性变换，最后使用softmax进行处理得到注意力张量
        attn_weights = F.softmax(self.attn(torch.cat((Q[0], K[0]), 1)), dim=1)
        # 将注意力矩阵和V进行一次bmm计算
        attn_applied = torch.bmm(attn_weights.unsqueeze(0), V)
        # 再次去Q[0]进行降维，再次和上面的运算结算进行拼接
        output = torch.cat((Q[0], attn_applied[0]), 1)
        # 第3步将上面的输出进行一次线性变换，然后再扩展维度成3维张量
        output = self.attn_combine(output).unsqueeze(0)
        return output, attn_weights


query_size = 32
key_size = 32
value_size1 = 32
value_size2 = 64
output_size = 64

attn = Attention(query_size, key_size, value_size1, value_size2, output_size)
Q = torch.randn(1, 1, 32)
K = torch.randn(1, 1, 32)
V = torch.randn(1, 32, 64)
output = attn(Q, K, V)
print(output[0])
print(output[0].size())
print(output[1])
print(output[1].size())
